{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rG2rwLmwh4Rc"
      },
      "source": [
        "#### Copyright 2023 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yOzhgsd6hz6X"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VfjRf_jVg43D"
      },
      "source": [
        "# YOLO Image classification\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "keCteDOniuMA"
      },
      "source": [
        "This tutorial trains [Darkent](https://arxiv.org/abs/1911.11929) from Tensorflow Model Garden package ([tf-models-official](https://pypi.org/project/tf-models-official/)) to classify images in the [cats_vs_dogs](https://www.tensorflow.org/datasets/catalog/cats_vs_dogs) dataset.\n",
        "\n",
        "\n",
        "[Model Garden](https://github.com/tensorflow/models/tree/master/official) contains a collection of state-of-the-art vision models, implemented with TensorFlow's high-level APIs. The implementations demonstrate the best practices for modeling, letting users to take full advantage of TensorFlow for their research and product development.\n",
        "\n",
        "**Dataset:** cats_vs_dogs\n",
        "* A large set of images of cats and dogs.\n",
        "\n",
        "This tutorial demonstrates how to:\n",
        "* Use models from the [TensorFlow Models package](https://pypi.org/project/tf-models-official/)\n",
        "* Train/Fine-tune a pre-built Darkent variations for Image Classification\n",
        "* Export the trained/tuned darknet model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "41irlZ-FlMhi"
      },
      "source": [
        "## Clone the model-garden repository"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "snEeWhVIdWgi"
      },
      "outputs": [],
      "source": [
        "! git clone -q https://github.com/tensorflow/models.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DVSglGuXmrjN"
      },
      "outputs": [],
      "source": [
        "! pip install -q -U tensorflow_datasets\n",
        "! pip install -q --user -r models/official/requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y_THKv-WwjmF"
      },
      "source": [
        "**Note:** Please restart runtime and continue with running the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XVQ1By43lc_k"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import sys\n",
        "\n",
        "import os\n",
        "os.environ['PYTHONPATH'] += \":/content/models\"\n",
        "\n",
        "sys.path.append(\"/content/models\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zEwBh_eplVZg"
      },
      "source": [
        "## Import necessary libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FTOX9q8tZE9R"
      },
      "outputs": [],
      "source": [
        "import pprint\n",
        "import logging\n",
        "import matplotlib.pyplot as plt\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "from official import core\n",
        "from official.vision.data import tfrecord_lib\n",
        "from official.vision import configs\n",
        "from official.vision.configs import common\n",
        "from official.projects.yolo.common import registry_imports\n",
        "from official.projects.yolo.serving import export_saved_model\n",
        "from official.projects.yolo.serving import export_module_factory\n",
        "from official.vision.serving import export_saved_model_lib\n",
        "\n",
        "logging.disable(logging.WARNING)\n",
        "pp = pprint.PrettyPrinter(indent=4)\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Obap4N9HlZzf"
      },
      "source": [
        "## Load dataset from Tensorflow Datasets(tfds)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 1551,
          "status": "ok",
          "timestamp": 1699997729595,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "_k5x1Tt8G9VV",
        "outputId": "e78c2377-65ec-4524-8c8f-25490a705256"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tfds.core.DatasetInfo(\n",
              "    name='cats_vs_dogs',\n",
              "    full_name='cats_vs_dogs/4.0.0',\n",
              "    description=\"\"\"\n",
              "    A large set of images of cats and dogs. There are 1738 corrupted images that are dropped.\n",
              "    \"\"\",\n",
              "    homepage='https://www.microsoft.com/en-us/download/details.aspx?id=54765',\n",
              "    data_dir='/root/tensorflow_datasets/cats_vs_dogs/4.0.0',\n",
              "    file_format=tfrecord,\n",
              "    download_size=786.67 MiB,\n",
              "    dataset_size=689.64 MiB,\n",
              "    features=FeaturesDict({\n",
              "        'image': Image(shape=(None, None, 3), dtype=uint8),\n",
              "        'image/filename': Text(shape=(), dtype=string),\n",
              "        'label': ClassLabel(shape=(), dtype=int64, num_classes=2),\n",
              "    }),\n",
              "    supervised_keys=('image', 'label'),\n",
              "    disable_shuffling=False,\n",
              "    splits={\n",
              "        'train': \u003cSplitInfo num_examples=23262, num_shards=8\u003e,\n",
              "    },\n",
              "    citation=\"\"\"@Inproceedings (Conference){asirra-a-captcha-that-exploits-interest-aligned-manual-image-categorization,\n",
              "    author = {Elson, Jeremy and Douceur, John (JD) and Howell, Jon and Saul, Jared},\n",
              "    title = {Asirra: A CAPTCHA that Exploits Interest-Aligned Manual Image Categorization},\n",
              "    booktitle = {Proceedings of 14th ACM Conference on Computer and Communications Security (CCS)},\n",
              "    year = {2007},\n",
              "    month = {October},\n",
              "    publisher = {Association for Computing Machinery, Inc.},\n",
              "    url = {https://www.microsoft.com/en-us/research/publication/asirra-a-captcha-that-exploits-interest-aligned-manual-image-categorization/},\n",
              "    edition = {Proceedings of 14th ACM Conference on Computer and Communications Security (CCS)},\n",
              "    }\"\"\",\n",
              ")"
            ]
          },
          "execution_count": 3,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "(train_ds, validation_ds, test_ds), ds_info = tfds.load(\n",
        "    name='cats_vs_dogs',\n",
        "    split=['train[:70%]', 'train[70%:90%]', 'train[90%:100%]'],\n",
        "    with_info=True)\n",
        "label_info = ds_info.features['label']\n",
        "ds_info"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "el_oHjBAlreD"
      },
      "source": [
        "## Write data to TFrecords"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qMp811pgl1_0"
      },
      "source": [
        "### Helper functions to preproces the data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JDIRTV1Bh4FA"
      },
      "outputs": [],
      "source": [
        "def process_record(record):\n",
        "    \"\"\"\n",
        "    Process a single record for TFRecords.\n",
        "\n",
        "    This function takes a record, typically containing image and label data,\n",
        "    and converts it into a TFRecord example. Detailed explaination is available here\n",
        "    https://www.tensorflow.org/api_docs/python/tf/train/Example\n",
        "\n",
        "    Args:\n",
        "        record (dict): A dictionary containing the record data with the following keys:\n",
        "            - 'image': A tensor representing the image data.\n",
        "            - 'label': A tensor representing the label associated with the image.\n",
        "\n",
        "    Returns:\n",
        "        tf.train.Example: A TFRecord example containing the processed data with\n",
        "        the following features:\n",
        "            - 'image/encoded': The encoded image data as a feature.\n",
        "            - 'image/class/label': The label data as a feature.\n",
        "    \"\"\"\n",
        "    keys_to_features = {\n",
        "        'image/encoded': tfrecord_lib.convert_to_feature(\n",
        "            tf.io.encode_jpeg(record['image']).numpy()),\n",
        "        'image/class/label': tfrecord_lib.convert_to_feature(\n",
        "            record['label'].numpy())\n",
        "    }\n",
        "    example = tf.train.Example(features=tf.train.Features(feature=keys_to_features))\n",
        "    return example\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pUQmR7seiCBF"
      },
      "outputs": [],
      "source": [
        "def write_tfrecords(dataset, output_path, num_shards=1):\n",
        "    \"\"\"\n",
        "    Write a dataset to TFRecords files.\n",
        "\n",
        "    This function takes a dataset and writes it to one or more TFRecords files,\n",
        "    splitting the data into shards if specified.\n",
        "\n",
        "    Args:\n",
        "        dataset (iterable): An iterable containing the data records to be written\n",
        "            to TFRecords. Each record should be in a format suitable for processing\n",
        "            with the 'process_record' function.\n",
        "        output_path (str): The base path where the TFRecords files will be saved.\n",
        "            If 'num_shards' is greater than 1, a unique suffix for each shard will\n",
        "            be added to the base path.\n",
        "        num_shards (int, optional): The number of TFRecords files to split the data\n",
        "            into. Defaults to 1, indicating no sharding.\n",
        "\n",
        "    Reuturns:\n",
        "        None\n",
        "    \"\"\"\n",
        "    writers = [\n",
        "        tf.io.TFRecordWriter(\n",
        "            output_path + '-%05d-of-%05d.tfrecord' % (i, num_shards))\n",
        "        for i in range(num_shards)\n",
        "    ]\n",
        "    for idx, record in enumerate(dataset):\n",
        "        if idx % LOG_EVERY == 0:\n",
        "            print('On image %d' % idx)\n",
        "        tf_example = process_record(record)\n",
        "        writers[idx % num_shards].write(tf_example.SerializeToString())\n",
        "\n",
        "LOG_EVERY = 1000\n",
        "output_dir = './cat_vs_dogs_tfrecords/'\n",
        "if not os.path.exists(output_dir):\n",
        "  os.mkdir(output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K35c-4M1ojdC"
      },
      "source": [
        "### Writing training data to TFRecords"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 23362,
          "status": "ok",
          "timestamp": 1699997756512,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "xiya20KViEBx",
        "outputId": "ab75364e-be3f-46f8-d0b1-e3186b0906d1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "On image 0\n",
            "On image 1000\n",
            "On image 2000\n",
            "On image 3000\n",
            "On image 4000\n",
            "On image 5000\n",
            "On image 6000\n",
            "On image 7000\n",
            "On image 8000\n",
            "On image 9000\n",
            "On image 10000\n",
            "On image 11000\n",
            "On image 12000\n",
            "On image 13000\n",
            "On image 14000\n",
            "On image 15000\n",
            "On image 16000\n"
          ]
        }
      ],
      "source": [
        "output_train_tfrecs = output_dir + 'train'\n",
        "write_tfrecords(train_ds, output_train_tfrecs,\n",
        "                num_shards=int(train_ds.cardinality().numpy() * 0.1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oAOuGouiolhq"
      },
      "source": [
        "### Writing validation data to TFRecords"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 14426,
          "status": "ok",
          "timestamp": 1699997770935,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "BdU2-OlLiHNh",
        "outputId": "fcffd93d-f24f-407b-d1c6-26880b1ac690"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "On image 0\n",
            "On image 1000\n",
            "On image 2000\n",
            "On image 3000\n",
            "On image 4000\n"
          ]
        }
      ],
      "source": [
        "output_validation_tfrecs = output_dir + 'validation'\n",
        "write_tfrecords(validation_ds, output_validation_tfrecs,\n",
        "                num_shards=int(validation_ds.cardinality().numpy() *0.1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iAU3OLKjosAd"
      },
      "source": [
        "### Writing testing data to TFRecords"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 32,
          "status": "ok",
          "timestamp": 1699997770937,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "8bcmO71biJR4",
        "outputId": "8c8bf0d9-8661-42ed-9e70-8ba30eb1a88b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "On image 0\n",
            "On image 1000\n",
            "On image 2000\n"
          ]
        }
      ],
      "source": [
        "output_test_tfrecs = output_dir + 'test'\n",
        "write_tfrecords(test_ds, output_test_tfrecs,\n",
        "                num_shards=int(test_ds.cardinality().numpy() *0.1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lc6Wf_oMozOJ"
      },
      "source": [
        "## Experiment Configuration"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FL5aunZjpIl_"
      },
      "source": [
        "### Load the existing configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tbqB6L5wldIE"
      },
      "outputs": [],
      "source": [
        "exp_config = core.exp_factory.get_exp_config('darknet_classification')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "09iJzu27pNaf"
      },
      "source": [
        "### Change the configuration parameters for custom dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "38duTtEKiSdr"
      },
      "outputs": [],
      "source": [
        "BATCH_SIZE = 16\n",
        "IMG_SIZE = 224\n",
        "\n",
        "epochs = 10\n",
        "steps_per_epoch = int(train_ds.cardinality().numpy() / BATCH_SIZE)\n",
        "validation_steps =  int(validation_ds.cardinality().numpy() / BATCH_SIZE)\n",
        "num_steps = epochs * steps_per_epoch\n",
        "\n",
        "lr = 0.012\n",
        "warmpup_lr = 0.1 * lr\n",
        "\n",
        "exp_config.task.model.input_size = [IMG_SIZE, IMG_SIZE, 3]\n",
        "exp_config.task.model.num_classes = ds_info.features['label'].num_classes\n",
        "\n",
        "exp_config.task.train_data.input_path = f'{output_train_tfrecs}*'\n",
        "exp_config.task.train_data.global_batch_size = BATCH_SIZE\n",
        "\n",
        "exp_config.task.validation_data.input_path = f'{output_validation_tfrecs}*'\n",
        "exp_config.task.validation_data.global_batch_size = BATCH_SIZE\n",
        "\n",
        "exp_config.trainer.checkpoint_interval = steps_per_epoch\n",
        "exp_config.trainer.best_checkpoint_export_subdir = 'best_ckpt'\n",
        "exp_config.trainer.optimizer_config.optimizer.type = 'sgd'\n",
        "exp_config.trainer.optimizer_config.optimizer.sgd.momentum = 0.9\n",
        "exp_config.trainer.optimizer_config.learning_rate.type = 'cosine'\n",
        "exp_config.trainer.optimizer_config.learning_rate.cosine.decay_steps = num_steps\n",
        "exp_config.trainer.optimizer_config.learning_rate.cosine.initial_learning_rate = lr\n",
        "exp_config.trainer.optimizer_config.warmup.type = 'linear'\n",
        "exp_config.trainer.optimizer_config.warmup.linear.warmup_learning_rate = warmpup_lr\n",
        "exp_config.trainer.optimizer_config.warmup.linear.warmup_steps = int(0.1 * steps_per_epoch)\n",
        "\n",
        "exp_config.trainer.train_steps = num_steps\n",
        "exp_config.trainer.steps_per_loop = steps_per_epoch\n",
        "exp_config.trainer.validation_steps = validation_steps\n",
        "exp_config.trainer.validation_interval = steps_per_epoch\n",
        "exp_config.trainer.summary_interval = steps_per_epoch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FINslYHhqaPt"
      },
      "source": [
        "### Set up the distribution strategy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 30,
          "status": "ok",
          "timestamp": 1699997770939,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "3-e4HjyKqZfw",
        "outputId": "5977b054-ee4d-4e8e-ae18-55d903a5547f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Running on single GPU  /device:GPU:0\n",
            "Number of accelerators:  1\n"
          ]
        }
      ],
      "source": [
        "# Detect hardware\n",
        "try:\n",
        "  tpu_resolver = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n",
        "except ValueError:\n",
        "  tpu_resolver = None\n",
        "  gpus = tf.config.experimental.list_logical_devices(\"GPU\")\n",
        "\n",
        "# Select appropriate distribution strategy\n",
        "if tpu_resolver:\n",
        "  tf.config.experimental_connect_to_cluster(tpu_resolver)\n",
        "  tf.tpu.experimental.initialize_tpu_system(tpu_resolver)\n",
        "  distribution_strategy = tf.distribute.experimental.TPUStrategy(tpu_resolver)\n",
        "  print('Running on TPU ', tpu_resolver.cluster_spec().as_dict()['worker'])\n",
        "elif len(gpus) \u003e 1:\n",
        "  distribution_strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])\n",
        "  print('Running on multiple GPUs ', [gpu.name for gpu in gpus])\n",
        "elif len(gpus) == 1:\n",
        "  distribution_strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on single GPU ', gpus[0].name)\n",
        "else:\n",
        "  distribution_strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on CPU')\n",
        "\n",
        "print(\"Number of accelerators: \", distribution_strategy.num_replicas_in_sync)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ugiuA0-0qq4Z"
      },
      "source": [
        "### Check the new configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 29,
          "status": "ok",
          "timestamp": 1699997770940,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "a5CO8qTqHqGH",
        "outputId": "d9d97856-051f-457b-a9d7-bd0cb5135042"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'runtime': {'all_reduce_alg': None,\n",
            "             'batchnorm_spatial_persistent': False,\n",
            "             'dataset_num_private_threads': None,\n",
            "             'default_shard_dim': -1,\n",
            "             'distribution_strategy': 'mirrored',\n",
            "             'enable_xla': False,\n",
            "             'gpu_thread_mode': None,\n",
            "             'loss_scale': None,\n",
            "             'mixed_precision_dtype': None,\n",
            "             'num_cores_per_replica': 1,\n",
            "             'num_gpus': 0,\n",
            "             'num_packs': 1,\n",
            "             'per_gpu_thread_count': 0,\n",
            "             'run_eagerly': False,\n",
            "             'task_index': -1,\n",
            "             'tpu': None,\n",
            "             'tpu_enable_xla_dynamic_padder': None,\n",
            "             'use_tpu_mp_strategy': False,\n",
            "             'worker_hosts': None},\n",
            " 'task': {'allow_image_summary': False,\n",
            "          'differential_privacy_config': None,\n",
            "          'evaluation': {'precision_and_recall_thresholds': None,\n",
            "                         'report_per_class_precision_and_recall': False,\n",
            "                         'top_k': 5},\n",
            "          'freeze_backbone': False,\n",
            "          'gradient_clip_norm': 0.0,\n",
            "          'init_checkpoint': '',\n",
            "          'logging_dir': None,\n",
            "          'losses': {'l2_weight_decay': 0.0,\n",
            "                     'label_smoothing': 0.0,\n",
            "                     'loss_weight': 1.0,\n",
            "                     'one_hot': True,\n",
            "                     'soft_labels': False,\n",
            "                     'use_binary_cross_entropy': False},\n",
            "          'model': {'add_head_batch_norm': False,\n",
            "                    'backbone': {'darknet': {'depth_scale': 1.0,\n",
            "                                             'dilate': False,\n",
            "                                             'max_level': 5,\n",
            "                                             'min_level': 3,\n",
            "                                             'model_id': 'cspdarknet53',\n",
            "                                             'use_reorg_input': False,\n",
            "                                             'use_separable_conv': False,\n",
            "                                             'width_scale': 1.0},\n",
            "                                 'type': 'darknet'},\n",
            "                    'dropout_rate': 0.0,\n",
            "                    'input_size': [224, 224, 3],\n",
            "                    'kernel_initializer': 'VarianceScaling',\n",
            "                    'norm_activation': {'activation': 'relu',\n",
            "                                        'norm_epsilon': 0.001,\n",
            "                                        'norm_momentum': 0.99,\n",
            "                                        'use_sync_bn': True},\n",
            "                    'num_classes': 2},\n",
            "          'name': None,\n",
            "          'train_data': {'apply_tf_data_service_before_batching': False,\n",
            "                         'aug_crop': True,\n",
            "                         'aug_policy': None,\n",
            "                         'aug_rand_hflip': True,\n",
            "                         'aug_type': None,\n",
            "                         'autotune_algorithm': None,\n",
            "                         'block_length': 1,\n",
            "                         'cache': False,\n",
            "                         'center_crop_fraction': 0.875,\n",
            "                         'color_jitter': 0.0,\n",
            "                         'crop_area_range': (0.08, 1.0),\n",
            "                         'cycle_length': 10,\n",
            "                         'decode_jpeg_only': True,\n",
            "                         'decoder': {'simple_decoder': {'attribute_names': [],\n",
            "                                                        'mask_binarize_threshold': None,\n",
            "                                                        'regenerate_source_id': False},\n",
            "                                     'type': 'simple_decoder'},\n",
            "                         'deterministic': None,\n",
            "                         'drop_remainder': True,\n",
            "                         'dtype': 'float32',\n",
            "                         'enable_shared_tf_data_service_between_parallel_trainers': False,\n",
            "                         'enable_tf_data_service': False,\n",
            "                         'file_type': 'tfrecord',\n",
            "                         'global_batch_size': 16,\n",
            "                         'image_field_key': 'image/encoded',\n",
            "                         'input_path': './cat_vs_dogs_tfrecords/train*',\n",
            "                         'is_multilabel': False,\n",
            "                         'is_training': True,\n",
            "                         'label_field_key': 'image/class/label',\n",
            "                         'mixup_and_cutmix': None,\n",
            "                         'prefetch_buffer_size': None,\n",
            "                         'randaug_magnitude': 10,\n",
            "                         'random_erasing': None,\n",
            "                         'repeated_augment': None,\n",
            "                         'seed': None,\n",
            "                         'sharding': True,\n",
            "                         'shuffle_buffer_size': 10000,\n",
            "                         'tf_data_service_address': None,\n",
            "                         'tf_data_service_job_name': None,\n",
            "                         'tf_resize_method': 'bilinear',\n",
            "                         'tfds_as_supervised': False,\n",
            "                         'tfds_data_dir': '',\n",
            "                         'tfds_name': '',\n",
            "                         'tfds_skip_decoding_feature': '',\n",
            "                         'tfds_split': '',\n",
            "                         'three_augment': False,\n",
            "                         'trainer_id': None,\n",
            "                         'weights': None},\n",
            "          'validation_data': {'apply_tf_data_service_before_batching': False,\n",
            "                              'aug_crop': True,\n",
            "                              'aug_policy': None,\n",
            "                              'aug_rand_hflip': True,\n",
            "                              'aug_type': None,\n",
            "                              'autotune_algorithm': None,\n",
            "                              'block_length': 1,\n",
            "                              'cache': False,\n",
            "                              'center_crop_fraction': 0.875,\n",
            "                              'color_jitter': 0.0,\n",
            "                              'crop_area_range': (0.08, 1.0),\n",
            "                              'cycle_length': 10,\n",
            "                              'decode_jpeg_only': True,\n",
            "                              'decoder': {'simple_decoder': {'attribute_names': [],\n",
            "                                                             'mask_binarize_threshold': None,\n",
            "                                                             'regenerate_source_id': False},\n",
            "                                          'type': 'simple_decoder'},\n",
            "                              'deterministic': None,\n",
            "                              'drop_remainder': True,\n",
            "                              'dtype': 'float32',\n",
            "                              'enable_shared_tf_data_service_between_parallel_trainers': False,\n",
            "                              'enable_tf_data_service': False,\n",
            "                              'file_type': 'tfrecord',\n",
            "                              'global_batch_size': 16,\n",
            "                              'image_field_key': 'image/encoded',\n",
            "                              'input_path': './cat_vs_dogs_tfrecords/validation*',\n",
            "                              'is_multilabel': False,\n",
            "                              'is_training': False,\n",
            "                              'label_field_key': 'image/class/label',\n",
            "                              'mixup_and_cutmix': None,\n",
            "                              'prefetch_buffer_size': None,\n",
            "                              'randaug_magnitude': 10,\n",
            "                              'random_erasing': None,\n",
            "                              'repeated_augment': None,\n",
            "                              'seed': None,\n",
            "                              'sharding': True,\n",
            "                              'shuffle_buffer_size': 10000,\n",
            "                              'tf_data_service_address': None,\n",
            "                              'tf_data_service_job_name': None,\n",
            "                              'tf_resize_method': 'bilinear',\n",
            "                              'tfds_as_supervised': False,\n",
            "                              'tfds_data_dir': '',\n",
            "                              'tfds_name': '',\n",
            "                              'tfds_skip_decoding_feature': '',\n",
            "                              'tfds_split': '',\n",
            "                              'three_augment': False,\n",
            "                              'trainer_id': None,\n",
            "                              'weights': None}},\n",
            " 'trainer': {'allow_tpu_summary': False,\n",
            "             'best_checkpoint_eval_metric': '',\n",
            "             'best_checkpoint_export_subdir': 'best_ckpt',\n",
            "             'best_checkpoint_metric_comp': 'higher',\n",
            "             'checkpoint_interval': 1017,\n",
            "             'continuous_eval_timeout': 3600,\n",
            "             'eval_tf_function': True,\n",
            "             'eval_tf_while_loop': False,\n",
            "             'loss_upper_bound': 1000000.0,\n",
            "             'max_to_keep': 5,\n",
            "             'optimizer_config': {'ema': None,\n",
            "                                  'learning_rate': {'cosine': {'alpha': 0.0,\n",
            "                                                               'decay_steps': 10170,\n",
            "                                                               'initial_learning_rate': 0.012,\n",
            "                                                               'name': 'CosineDecay',\n",
            "                                                               'offset': 0},\n",
            "                                                    'type': 'cosine'},\n",
            "                                  'optimizer': {'sgd': {'clipnorm': None,\n",
            "                                                        'clipvalue': None,\n",
            "                                                        'decay': 0.0,\n",
            "                                                        'global_clipnorm': None,\n",
            "                                                        'momentum': 0.9,\n",
            "                                                        'name': 'SGD',\n",
            "                                                        'nesterov': False},\n",
            "                                                'type': 'sgd'},\n",
            "                                  'warmup': {'linear': {'name': 'linear',\n",
            "                                                        'warmup_learning_rate': 0.0012000000000000001,\n",
            "                                                        'warmup_steps': 101},\n",
            "                                             'type': 'linear'}},\n",
            "             'preemption_on_demand_checkpoint': True,\n",
            "             'recovery_begin_steps': 0,\n",
            "             'recovery_max_trials': 0,\n",
            "             'steps_per_loop': 1017,\n",
            "             'summary_interval': 1017,\n",
            "             'train_steps': 10170,\n",
            "             'train_tf_function': True,\n",
            "             'train_tf_while_loop': True,\n",
            "             'validation_interval': 1017,\n",
            "             'validation_steps': 290,\n",
            "             'validation_summary_subdir': 'validation'}}\n"
          ]
        }
      ],
      "source": [
        "pprint.pprint(exp_config.as_dict())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s21NhXxwqzcD"
      },
      "source": [
        "## Create the Task object `(tfm.core.base_task.Task)` from the `config_definitions.TaskConfig`.\n",
        "\n",
        "The Task object has all the methods necessary for **building the dataset, building the model, and running training \u0026 evaluation**. These methods are driven by `tfm.core.train_lib.run_experiment`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n6TYLP2mIh7H"
      },
      "outputs": [],
      "source": [
        "model_dir = './trained_model/'\n",
        "with distribution_strategy.scope():\n",
        "  task = core.task_factory.get_task(exp_config.task, logging_dir=model_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 29,
          "status": "ok",
          "timestamp": 1699997770943,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "cdC3Plx7rMTi",
        "outputId": "bf02a8e2-3446-4c09-e15d-373fa928bc6b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "images.shape: (16, 224, 224, 3)  images.dtype: tf.float32\n",
            "labels.shape: (16,)             labels.dtype: tf.int32\n"
          ]
        }
      ],
      "source": [
        "for images, labels in task.build_inputs(exp_config.task.train_data).take(1):\n",
        "  print(f'images.shape: {str(images.shape):16}  images.dtype: {images.dtype!r}')\n",
        "  print(f'labels.shape: {str(labels.shape):16}  labels.dtype: {labels.dtype!r}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0-KFVW4KA_vU"
      },
      "source": [
        "## Save the configuration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "onc4YJbBBE3e"
      },
      "outputs": [],
      "source": [
        "core.train_utils.serialize_config(exp_config, model_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KasKyG_hrYac"
      },
      "source": [
        "## Visualize the training data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "emn7QfE1rfQc"
      },
      "source": [
        "### Use `ds_info` (which is an instance of tfds.core.DatasetInfo) to lookup the text descriptions of each class ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MmZg2Z3drToR"
      },
      "outputs": [],
      "source": [
        "label_info = ds_info.features['label']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3_xSuXPrnX_"
      },
      "source": [
        "### Visualize a batch of the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CQBTmQpzrrGb"
      },
      "outputs": [],
      "source": [
        "def show_batch(images, labels, predictions=None):\n",
        "  plt.figure(figsize=(10, 10))\n",
        "  min = images.numpy().min()\n",
        "  max = images.numpy().max()\n",
        "  delta = max - min\n",
        "\n",
        "  for i in range(BATCH_SIZE):\n",
        "    plt.subplot(4, 4, i + 1)\n",
        "    plt.imshow((images[i]-min) / delta)\n",
        "    if predictions is None:\n",
        "      plt.title(label_info.int2str(labels[i]))\n",
        "    else:\n",
        "      if labels[i] == predictions[i]:\n",
        "        color = 'g'\n",
        "      else:\n",
        "        color = 'r'\n",
        "      plt.title(label_info.int2str(predictions[i]), color=color)\n",
        "    plt.axis(\"off\")\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Du8GwFayrydM"
      },
      "outputs": [],
      "source": [
        "for images, labels in task.build_inputs(exp_config.task.validation_data).take(1):\n",
        "  show_batch(images, labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oz9M3EDZro3W"
      },
      "source": [
        "## Train and Evaluate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 1214820,
          "status": "ok",
          "timestamp": 1699997521229,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "492B5k0cIo0e",
        "outputId": "cb486eb5-5504-4e6c-e46d-df1e807d83e1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "restoring or initializing model...\n",
            "train | step:      0 | training until step 1017...\n",
            "train | step:   1017 | steps/sec:    6.4 | output: \n",
            "    {'accuracy': 0.54013026,\n",
            "     'learning_rate': 0.011706339,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.7053006}\n",
            "saved checkpoint to ./trained_model/ckpt-1017.\n",
            " eval | step:   1017 | running 290 steps of evaluation...\n",
            " eval | step:   1017 | steps/sec:   17.6 | eval time:   16.4 sec | output: \n",
            "    {'accuracy': 0.57693964,\n",
            "     'steps_per_second': 17.634644012635317,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.67286325}\n",
            "train | step:   1017 | training until step 2034...\n",
            "train | step:   2034 | steps/sec:    8.6 | output: \n",
            "    {'accuracy': 0.58898723,\n",
            "     'learning_rate': 0.010854102,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.6713662}\n",
            "saved checkpoint to ./trained_model/ckpt-2034.\n",
            " eval | step:   2034 | running 290 steps of evaluation...\n",
            " eval | step:   2034 | steps/sec:   37.3 | eval time:    7.8 sec | output: \n",
            "    {'accuracy': 0.6153017,\n",
            "     'steps_per_second': 37.2743481981908,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.652928}\n",
            "train | step:   2034 | training until step 3051...\n",
            "train | step:   3051 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.6157817,\n",
            "     'learning_rate': 0.009526712,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.65529084}\n",
            "saved checkpoint to ./trained_model/ckpt-3051.\n",
            " eval | step:   3051 | running 290 steps of evaluation...\n",
            " eval | step:   3051 | steps/sec:   37.2 | eval time:    7.8 sec | output: \n",
            "    {'accuracy': 0.51681036,\n",
            "     'steps_per_second': 37.18629141886086,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 1.0405935}\n",
            "train | step:   3051 | training until step 4068...\n",
            "train | step:   4068 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.6411013,\n",
            "     'learning_rate': 0.007854102,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.63446975}\n",
            "saved checkpoint to ./trained_model/ckpt-4068.\n",
            " eval | step:   4068 | running 290 steps of evaluation...\n",
            " eval | step:   4068 | steps/sec:   37.3 | eval time:    7.8 sec | output: \n",
            "    {'accuracy': 0.68297416,\n",
            "     'steps_per_second': 37.306989215139936,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.59876406}\n",
            "train | step:   4068 | training until step 5085...\n",
            "train | step:   5085 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.6549287,\n",
            "     'learning_rate': 0.0059999996,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.6212429}\n",
            "saved checkpoint to ./trained_model/ckpt-5085.\n",
            " eval | step:   5085 | running 290 steps of evaluation...\n",
            " eval | step:   5085 | steps/sec:   37.7 | eval time:    7.7 sec | output: \n",
            "    {'accuracy': 0.67349136,\n",
            "     'steps_per_second': 37.65721170518001,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.59417003}\n",
            "train | step:   5085 | training until step 6102...\n",
            "train | step:   6102 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.6778515,\n",
            "     'learning_rate': 0.004145897,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.5970689}\n",
            "saved checkpoint to ./trained_model/ckpt-6102.\n",
            " eval | step:   6102 | running 290 steps of evaluation...\n",
            " eval | step:   6102 | steps/sec:   37.3 | eval time:    7.8 sec | output: \n",
            "    {'accuracy': 0.68168104,\n",
            "     'steps_per_second': 37.324507959229,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.6371893}\n",
            "train | step:   6102 | training until step 7119...\n",
            "train | step:   7119 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.7034169,\n",
            "     'learning_rate': 0.0024732884,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.56852204}\n",
            "saved checkpoint to ./trained_model/ckpt-7119.\n",
            " eval | step:   7119 | running 290 steps of evaluation...\n",
            " eval | step:   7119 | steps/sec:   37.4 | eval time:    7.7 sec | output: \n",
            "    {'accuracy': 0.7521552,\n",
            "     'steps_per_second': 37.433325893948016,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.52360195}\n",
            "train | step:   7119 | training until step 8136...\n",
            "train | step:   8136 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.7234513,\n",
            "     'learning_rate': 0.0011458977,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.54219157}\n",
            "saved checkpoint to ./trained_model/ckpt-8136.\n",
            " eval | step:   8136 | running 290 steps of evaluation...\n",
            " eval | step:   8136 | steps/sec:   37.5 | eval time:    7.7 sec | output: \n",
            "    {'accuracy': 0.787931,\n",
            "     'steps_per_second': 37.471471146938974,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.45562845}\n",
            "train | step:   8136 | training until step 9153...\n",
            "train | step:   9153 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.7457596,\n",
            "     'learning_rate': 0.0002936611,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.50991994}\n",
            "saved checkpoint to ./trained_model/ckpt-9153.\n",
            " eval | step:   9153 | running 290 steps of evaluation...\n",
            " eval | step:   9153 | steps/sec:   37.7 | eval time:    7.7 sec | output: \n",
            "    {'accuracy': 0.80625,\n",
            "     'steps_per_second': 37.66015917575601,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.41287535}\n",
            "train | step:   9153 | training until step 10170...\n",
            "train | step:  10170 | steps/sec:    9.3 | output: \n",
            "    {'accuracy': 0.7517822,\n",
            "     'learning_rate': 0.0,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'training_loss': 0.5036062}\n",
            "saved checkpoint to ./trained_model/ckpt-10170.\n",
            " eval | step:  10170 | running 290 steps of evaluation...\n",
            " eval | step:  10170 | steps/sec:   37.6 | eval time:    7.7 sec | output: \n",
            "    {'accuracy': 0.81314653,\n",
            "     'steps_per_second': 37.58710303226575,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.40270108}\n",
            " eval | step:  10170 | running 290 steps of evaluation...\n",
            " eval | step:  10170 | steps/sec:   37.3 | eval time:    7.8 sec | output: \n",
            "    {'accuracy': 0.81314653,\n",
            "     'steps_per_second': 37.3229583933457,\n",
            "     'top_5_accuracy': 1.0,\n",
            "     'validation_loss': 0.40270108}\n"
          ]
        }
      ],
      "source": [
        "model, eval_logs = core.train_lib.run_experiment(\n",
        "    distribution_strategy=distribution_strategy,\n",
        "    task=task,\n",
        "    mode='train_and_eval',\n",
        "    params=exp_config,\n",
        "    model_dir=model_dir,\n",
        "    run_post_eval=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yz8-qUIcFDKt"
      },
      "source": [
        "## Export the trained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 99764,
          "status": "ok",
          "timestamp": 1699997870693,
          "user": {
            "displayName": "Siva Sravana Kumar Neeli",
            "userId": "06669604936988620923"
          },
          "user_tz": 480
        },
        "id": "SDX0RWMBtZB2",
        "outputId": "c0a78a30-b278-4266-b8c2-41a68dd84cd0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "2023-11-14 21:36:10.722529: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
            "2023-11-14 21:36:18.463386: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:47] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.\n",
            "I1114 21:36:49.038166 140374729912320 signature_serialization.py:148] Function `serve` contains input name(s) resource with unsupported characters which will be renamed to dense_biasadd_readvariableop_resource in the SavedModel.\n",
            "I1114 21:37:34.619264 140374729912320 save.py:274] Found untraced functions such as serve_eval, conv2d_layer_call_fn, conv2d_layer_call_and_return_conditional_losses, _jit_compiled_convolution_op, conv_bn_layer_call_fn while saving (showing 5 of 415). These functions will not be directly callable after loading.\n",
            "INFO:tensorflow:Assets written to: ./exported_model//saved_model/assets\n",
            "I1114 21:37:46.884536 140374729912320 builder_impl.py:804] Assets written to: ./exported_model//saved_model/assets\n",
            "I1114 21:37:47.379255 140374729912320 fingerprinting_utils.py:48] Writing fingerprint to ./exported_model//saved_model/fingerprint.pb\n",
            "I1114 21:37:48.188559 140374729912320 train_utils.py:400] Saving experiment configuration to ./exported_model//params.yaml\n"
          ]
        }
      ],
      "source": [
        "EXPORT_DIR_PATH = \"./exported_model/\"\n",
        "!python -m official.projects.yolo.serving.export_saved_model \\\n",
        "  --experiment=\"darknet_classification\" \\\n",
        "  --export_dir=$EXPORT_DIR_PATH/ \\\n",
        "  --checkpoint_path=$model_dir \\\n",
        "  --config_file=$model_dir/params.yaml \\\n",
        "  --batch_size=$BATCH_SIZE \\\n",
        "  --input_type=\"image_tensor\" \\\n",
        "  --input_image_size=$IMG_SIZE,$IMG_SIZE"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h64DRrJEr7ls"
      },
      "source": [
        "### Test the exported model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JKcuFx17FBeu"
      },
      "source": [
        "### Importing SavedModel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2-GADU74zkjK"
      },
      "outputs": [],
      "source": [
        "imported = tf.saved_model.load('/content/exported_model/saved_model')\n",
        "model_fn = imported.signatures['serving_default']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNKz-BD-sFHJ"
      },
      "source": [
        "### Visualize the test predictions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9lE9sOYGxIkf"
      },
      "outputs": [],
      "source": [
        "def resize_image(record):\n",
        "  image = tf.image.resize(record['image'], size=(IMG_SIZE, IMG_SIZE))\n",
        "  image = tf.cast(image, tf.uint8)\n",
        "  return image, record['label']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81gCeHCava1w"
      },
      "outputs": [],
      "source": [
        "test_ds_resized = test_ds.map(resize_image).shuffle(100)\n",
        "test_ds_batched = test_ds_resized.batch(BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P7b5ApE2bRFW"
      },
      "outputs": [],
      "source": [
        "for images, labels in test_ds_batched.take(1):\n",
        "  predictions = model_fn(inputs=images)['logits']\n",
        "  predictions = tf.argmax(predictions, axis=-1)\n",
        "\n",
        "show_batch(images, labels, predictions)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [
        {
          "file_id": "1rt6hpJWOSKa2c746Ss58VERmnaA3gWys",
          "timestamp": 1700078413508
        }
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
